from pinecone import Pinecone, ServerlessSpec
from sklearn.datasets import load_digits  # 新增：用于加载数据集

# 初始化Pinecone客户端
pinecone = Pinecone(api_key="pcsk_xCptY_QRJADCio5ZJTrGTNgEsJHZrHA8jWLnvUmMehdDcQnX4wbR9Pk2CXH28NnkjyuWK")

# 索引名称
index_name = "mnist-index"

################### 检查、清除之前创建的索引 ###################
existing_indexes = pinecone.list_indexes()

if any(index['name'] == index_name for index in existing_indexes):
    print(f"索引 '{index_name}' 已存在，正在删除...")
    pinecone.delete_index(index_name)
    print(f"索引 '{index_name}' 已成功删除。")
else:
    print(f"索引 '{index_name}' 不存在，将创建新索引。")
################### 检查、清除之前创建的索引 ###################

# 创建新索引
print(f"正在创建新索引 '{index_name}'...")
pinecone.create_index(
    name=index_name,
    dimension=64,  # 8x8图像展平后的维度
    metric="euclidean",
    spec=ServerlessSpec(
        cloud="aws",
        region="us-east-1"
    )
)
print(f"索引 '{index_name}' 创建成功。")

# 连接到索引
index = pinecone.Index(index_name)
print(f"已成功连接到索引 '{index_name}'。")

# ---------------------- 新增：向量上传代码 ----------------------
# 加载MNIST数据集（8x8版本）
digits = load_digits(n_class=10)
X = digits.data  # 特征数据（64维向量）
y = digits.target  # 标签数据（0-9）

# 转换数据为Pinecone可接受的格式
vectors = []
for i in range(len(X)):
    vector_id = str(i)  # 用索引作为唯一ID
    vector_values = X[i].tolist()  # 转换为Python列表
    metadata = {"label": int(y[i])}  # 存储真实标签
    vectors.append((vector_id, vector_values, metadata))

# 分批上传向量（避免请求过大）
batch_size = 1000
for i in range(0, len(vectors), batch_size):
    batch = vectors[i:i + batch_size]
    index.upsert(batch)
    print(f"已上传 {min(i + batch_size, len(vectors))}/{len(vectors)} 个向量")

print("所有向量上传完成！")